import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, r2_score
import warnings
warnings.filterwarnings('ignore')

# 数据模拟程序
class MembraneFoulingDataGenerator:
    """膜污染模拟数据生成器"""

    def __init__(self, n_samples=1000, random_state=42):
        self.n_samples = n_samples
        self.random_state = random_state
        np.random.seed(random_state)

    def generate_operating_parameters(self):
        """生成运行参数"""
        data = {}

        # TMP: 膜前后压力 (bar)，随着污染增加而增加
        base_tmp = np.random.normal(1.5, 0.3, self.n_samples)
        time_factor = np.linspace(0, 1, self.n_samples)
        data['TMP'] = base_tmp + 0.8 * time_factor + 0.2 * np.random.normal(0, 0.1, self.n_samples)

        # dTMP_dt: 压差变化率 (bar/h)
        data['dTMP_dt'] = np.gradient(data['TMP']) * 10 + np.random.normal(0, 0.05, self.n_samples)

        # Q: 进水流量 (m³/h)
        data['Q'] = np.random.normal(50, 10, self.n_samples)

        # T: 温度 (°C)
        data['T'] = np.random.normal(25, 5, self.n_samples)

        return pd.DataFrame(data)

    def generate_water_quality_parameters(self):
        """生成水质参数"""
        data = {}

        # 进水水质
        data['TOC_in'] = np.random.lognormal(2.5, 0.3, self.n_samples)  # mg/L
        data['CODMn_in'] = np.random.lognormal(3.0, 0.4, self.n_samples)  # mg/L
        data['NTU_in'] = np.random.lognormal(1.2, 0.5, self.n_samples)  # NTU
        data['TDS_in'] = np.random.normal(800, 200, self.n_samples)  # mg/L

        # 产水水质 (通常比进水好)
        data['TOC_permeate'] = data['TOC_in'] * np.random.uniform(0.05, 0.2, self.n_samples)
        data['CODMn_permeate'] = data['CODMn_in'] * np.random.uniform(0.08, 0.25, self.n_samples)
        data['NTU_permeate'] = data['NTU_in'] * np.random.uniform(0.01, 0.1, self.n_samples)
        data['TDS_permeate'] = data['TDS_in'] * np.random.uniform(0.1, 0.3, self.n_samples)

        # 回水/废水水质
        data['TOC_concentrate'] = data['TOC_in'] * np.random.uniform(1.5, 3.0, self.n_samples)
        data['CODMn_concentrate'] = data['CODMn_in'] * np.random.uniform(1.6, 3.2, self.n_samples)
        data['NTU_concentrate'] = data['NTU_in'] * np.random.uniform(2.0, 4.0, self.n_samples)
        data['TDS_concentrate'] = data['TDS_in'] * np.random.uniform(1.8, 3.5, self.n_samples)

        return pd.DataFrame(data)

    def calculate_fouling_index(self, operating_params, quality_params):
        """计算膜污染指数"""
        # 基于运行参数和水质参数计算综合污染指数
        fouling_index = (
            0.3 * (operating_params['TMP'] - 1.0) +  # TMP贡献
            0.2 * np.abs(operating_params['dTMP_dt']) +  # 压差变化率贡献
            0.1 * (quality_params['TOC_in'] / 10) +  # TOC贡献
            0.1 * (quality_params['CODMn_in'] / 20) +  # CODMn贡献
            0.1 * (quality_params['NTU_in'] / 5) +  # 浊度贡献
            0.1 * (quality_params['TDS_in'] / 1000) +  # TDS贡献
            0.1 * ((quality_params['TOC_concentrate'] / quality_params['TOC_in']) - 1)  # 浓缩倍数
        )

        # 添加时间趋势
        time_factor = np.linspace(0, 1, self.n_samples)
        fouling_index += 0.5 * time_factor

        # 添加随机噪声
        fouling_index += np.random.normal(0, 0.1, self.n_samples)

        return np.clip(fouling_index, 0, 10)  # 限制在0-10范围内

    def calculate_remaining_life(self, fouling_index):
        """计算剩余使用寿命"""
        # 假设污染指数达到8时需要更换
        critical_fouling = 8.0
        max_life = 365  # 天

        remaining_life = max_life * (1 - fouling_index / critical_fouling)
        return np.clip(remaining_life, 0, max_life)

    def generate_complete_dataset(self):
        """生成完整数据集"""
        print("生成运行参数...")
        operating_params = self.generate_operating_parameters()

        print("生成水质参数...")
        quality_params = self.generate_water_quality_parameters()

        print("计算污染指数...")
        fouling_index = self.calculate_fouling_index(operating_params, quality_params)

        print("计算剩余寿命...")
        remaining_life = self.calculate_remaining_life(fouling_index)

        # 合并所有数据
        complete_data = pd.concat([operating_params, quality_params], axis=1)
        complete_data['Fouling_Index'] = fouling_index
        complete_data['Remaining_Life_Days'] = remaining_life
        complete_data['Operation_Days'] = np.linspace(0, 365, self.n_samples)

        return complete_data

# 生成模拟数据
print("开始生成膜污染模拟数据...")
generator = MembraneFoulingDataGenerator(n_samples=2000)
membrane_data = generator.generate_complete_dataset()

print(f"生成数据完成！共 {len(membrane_data)} 条记录")
print("\n数据前5行：")
print(membrane_data.head())
print("\n数据统计信息：")
print(membrane_data.describe())




# 膜污染评估程序
class MembraneFoulingEvaluator:
    """膜污染评估器"""

    def __init__(self):
        self.models = {}
        self.scalers = {}
        self.feature_importance = {}

    def prepare_features(self, data):
        """准备特征数据"""
        # 运行参数特征
        operating_features = ['TMP', 'dTMP_dt', 'Q', 'T']

        # 水质参数特征
        quality_features = [
            'TOC_in', 'CODMn_in', 'NTU_in', 'TDS_in',
            'TOC_permeate', 'CODMn_permeate', 'NTU_permeate', 'TDS_permeate',
            'TOC_concentrate', 'CODMn_concentrate', 'NTU_concentrate', 'TDS_concentrate'
        ]

        # 衍生特征
        data['TOC_rejection'] = (data['TOC_in'] - data['TOC_permeate']) / data['TOC_in']
        data['CODMn_rejection'] = (data['CODMn_in'] - data['CODMn_permeate']) / data['CODMn_in']
        data['NTU_rejection'] = (data['NTU_in'] - data['NTU_permeate']) / data['NTU_in']
        data['TDS_rejection'] = (data['TDS_in'] - data['TDS_permeate']) / data['TDS_in']

        derived_features = ['TOC_rejection', 'CODMn_rejection', 'NTU_rejection', 'TDS_rejection']

        all_features = operating_features + quality_features + derived_features

        return data[all_features], all_features

    def train_fouling_index_model(self, X, y):
        """训练污染指数预测模型"""
        # 数据标准化
        self.scalers['fouling'] = StandardScaler()
        X_scaled = self.scalers['fouling'].fit_transform(X)

        # 使用随机森林回归
        self.models['fouling'] = RandomForestRegressor(
            n_estimators=100,
            max_depth=10,
            random_state=42
        )

        X_train, X_test, y_train, y_test = train_test_split(
            X_scaled, y, test_size=0.2, random_state=42
        )

        self.models['fouling'].fit(X_train, y_train)

        # 评估模型
        y_pred = self.models['fouling'].predict(X_test)
        mse = mean_squared_error(y_test, y_pred)
        r2 = r2_score(y_test, y_pred)

        # 特征重要性
        self.feature_importance['fouling'] = dict(zip(
            self.feature_names,
            self.models['fouling'].feature_importances_
        ))

        return mse, r2

    def train_life_prediction_model(self, X, y):
        """训练剩余寿命预测模型"""
        # 数据标准化
        self.scalers['life'] = StandardScaler()
        X_scaled = self.scalers['life'].fit_transform(X)

        # 使用随机森林回归
        self.models['life'] = RandomForestRegressor(
            n_estimators=100,
            max_depth=10,
            random_state=42
        )

        X_train, X_test, y_train, y_test = train_test_split(
            X_scaled, y, test_size=0.2, random_state=42
        )

        self.models['life'].fit(X_train, y_train)

        # 评估模型
        y_pred = self.models['life'].predict(X_test)
        mse = mean_squared_error(y_test, y_pred)
        r2 = r2_score(y_test, y_pred)

        # 特征重要性
        self.feature_importance['life'] = dict(zip(
            self.feature_names,
            self.models['life'].feature_importances_
        ))

        return mse, r2

    def train_models(self, data):
        """训练所有模型"""
        print("准备特征数据...")
        X, self.feature_names = self.prepare_features(data)
        y_fouling = data['Fouling_Index']
        y_life = data['Remaining_Life_Days']

        print("训练污染指数预测模型...")
        mse_fouling, r2_fouling = self.train_fouling_index_model(X, y_fouling)

        print("训练剩余寿命预测模型...")
        mse_life, r2_life = self.train_life_prediction_model(X, y_life)

        print("\n=== 模型训练结果 ===")
        print(f"污染指数预测 - MSE: {mse_fouling:.4f}, R²: {r2_fouling:.4f}")
        print(f"剩余寿命预测 - MSE: {mse_life:.4f}, R²: {r2_life:.4f}")

        return mse_fouling, r2_fouling, mse_life, r2_life

    def predict_fouling(self, new_data):
        """预测膜污染指数"""
        X, _ = self.prepare_features(new_data)
        X_scaled = self.scalers['fouling'].transform(X)
        return self.models['fouling'].predict(X_scaled)

    def predict_remaining_life(self, new_data):
        """预测剩余寿命"""
        X, _ = self.prepare_features(new_data)
        X_scaled = self.scalers['life'].transform(X)
        return self.models['life'].predict(X_scaled)

    def plot_feature_importance(self):
        """绘制特征重要性图"""
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

        # 污染指数特征重要性
        fouling_importance = sorted(
            self.feature_importance['fouling'].items(),
            key=lambda x: x[1],
            reverse=True
        )[:10]

        features, importance = zip(*fouling_importance)
        ax1.barh(range(len(features)), importance)
        ax1.set_yticks(range(len(features)))
        ax1.set_yticklabels(features)
        ax1.set_xlabel('特征重要性')
        ax1.set_title('膜污染指数预测特征重要性')

        # 剩余寿命特征重要性
        life_importance = sorted(
            self.feature_importance['life'].items(),
            key=lambda x: x[1],
            reverse=True
        )[:10]

        features, importance = zip(*life_importance)
        ax2.barh(range(len(features)), importance)
        ax2.set_yticks(range(len(features)))
        ax2.set_yticklabels(features)
        ax2.set_xlabel('特征重要性')
        ax2.set_title('剩余寿命预测特征重要性')

        plt.tight_layout()
        plt.show()

    def plot_trend_analysis(self, data):
        """绘制趋势分析图"""
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))

        # TMP趋势
        ax1.scatter(data['Operation_Days'], data['TMP'], alpha=0.6, s=10)
        ax1.set_xlabel('运行天数')
        ax1.set_ylabel('TMP (bar)')
        ax1.set_title('膜前后压力趋势')
        ax1.grid(True, alpha=0.3)

        # 污染指数趋势
        ax2.scatter(data['Operation_Days'], data['Fouling_Index'], alpha=0.6, s=10)
        ax2.axhline(y=8, color='r', linestyle='--', label='临界污染指数')
        ax2.set_xlabel('运行天数')
        ax2.set_ylabel('污染指数')
        ax2.set_title('膜污染指数趋势')
        ax2.legend()
        ax2.grid(True, alpha=0.3)

        # 剩余寿命趋势
        ax3.scatter(data['Operation_Days'], data['Remaining_Life_Days'], alpha=0.6, s=10)
        ax3.set_xlabel('运行天数')
        ax3.set_ylabel('剩余寿命 (天)')
        ax3.set_title('膜剩余使用寿命趋势')
        ax3.grid(True, alpha=0.3)

        # TOC去除率趋势
        toc_rejection = (data['TOC_in'] - data['TOC_permeate']) / data['TOC_in']
        ax4.scatter(data['Operation_Days'], toc_rejection, alpha=0.6, s=10)
        ax4.set_xlabel('运行天数')
        ax4.set_ylabel('TOC去除率')
        ax4.set_title('TOC去除率趋势')
        ax4.grid(True, alpha=0.3)

        plt.tight_layout()
        plt.show()


# 训练评估模型
print("\n开始训练膜污染评估模型...")
evaluator = MembraneFoulingEvaluator()
mse_fouling, r2_fouling, mse_life, r2_life = evaluator.train_models(membrane_data)

# 可视化分析
print("\n生成可视化分析...")
evaluator.plot_feature_importance()
evaluator.plot_trend_analysis(membrane_data)

# 测试预测
print("\n测试预测功能...")
test_data = membrane_data.iloc[:10].copy()
fouling_predictions = evaluator.predict_fouling(test_data)
life_predictions = evaluator.predict_remaining_life(test_data)

print("测试预测结果:")
for i in range(3):  # 显示前3个样本
    print(f"样本 {i + 1}: 污染指数预测={fouling_predictions[i]:.2f}, 实际={test_data['Fouling_Index'].iloc[i]:.2f}, "
          f"剩余寿命预测={life_predictions[i]:.0f}天, 实际={test_data['Remaining_Life_Days'].iloc[i]:.0f}天")




# 实时监控和预警系统

class MembraneMonitoringSystem:
    """膜系统实时监控和预警系统"""

    def __init__(self, evaluator, warning_threshold=6, critical_threshold=8):
        self.evaluator = evaluator
        self.warning_threshold = warning_threshold
        self.critical_threshold = critical_threshold
        self.history_data = []

    def add_new_data(self, new_data):
        """添加新数据并进行分析"""
        # 预测污染指数和剩余寿命
        fouling_pred = self.evaluator.predict_fouling(new_data)[0]
        life_pred = self.evaluator.predict_remaining_life(new_data)[0]

        # 状态评估
        if fouling_pred >= self.critical_threshold:
            status = "CRITICAL"
            action = "立即进行膜清洗或更换"
        elif fouling_pred >= self.warning_threshold:
            status = "WARNING"
            action = "建议安排膜清洗"
        else:
            status = "NORMAL"
            action = "继续监控"

        result = {
            'fouling_index': fouling_pred,
            'remaining_life': life_pred,
            'status': status,
            'recommended_action': action,
            'timestamp': pd.Timestamp.now()
        }

        self.history_data.append(result)
        return result

    def generate_report(self, recent_data):
        """生成监控报告"""
        current_status = self.add_new_data(recent_data)

        print("\n" + "=" * 50)
        print("       膜系统污染状态监控报告")
        print("=" * 50)
        print(f"监测时间: {current_status['timestamp']}")
        print(f"当前污染指数: {current_status['fouling_index']:.2f}")
        print(f"预测剩余寿命: {current_status['remaining_life']:.0f} 天")
        print(f"系统状态: {current_status['status']}")
        print(f"建议措施: {current_status['recommended_action']}")
        print("=" * 50)

        # 显示关键运行参数
        print("\n关键运行参数:")
        print(f"  TMP: {recent_data['TMP'].iloc[0]:.2f} bar")
        print(f"  dTMP/dt: {recent_data['dTMP_dt'].iloc[0]:.3f} bar/h")
        print(f"  进水流量: {recent_data['Q'].iloc[0]:.1f} m³/h")
        print(f"  温度: {recent_data['T'].iloc[0]:.1f} °C")

        # 显示关键水质参数
        print("\n关键水质参数:")
        print(f"  进水TOC: {recent_data['TOC_in'].iloc[0]:.1f} mg/L")
        print(f"  产水TOC: {recent_data['TOC_permeate'].iloc[0]:.2f} mg/L")
        print(
            f"  TOC去除率: {(recent_data['TOC_in'].iloc[0] - recent_data['TOC_permeate'].iloc[0]) / recent_data['TOC_in'].iloc[0] * 100:.1f}%")


# 创建监控系统
monitoring_system = MembraneMonitoringSystem(evaluator)

# 模拟实时监控
print("\n模拟实时监控...")
for i in range(5):
    # 从数据中随机选择样本模拟新数据
    sample_idx = np.random.randint(100, 2000)
    sample_data = membrane_data.iloc[[sample_idx]]

    print(f"\n--- 监控周期 {i + 1} ---")
    monitoring_system.generate_report(sample_data)